#!/usr/bin/env python3
import os
import cv2
import time
from opt_visualizer import MergedPosePositionVisualizer
from shared_memory_reader import SharedMemoryReader
from seg_pose import SegPose
from compute_position import ComputePosition  
from utils import get_depth_image
from logger import setup_logger
import glob

logger = setup_logger()

def main():
    """主函数"""
    seg_pose = SegPose()
    compute_position = ComputePosition()
    st_img = cv2.imread("./data/Color/st.jpg")
    st_depth_img = get_depth_image("./data/Color/st.jpg")

    logger.info("==========设置标准位置==========")
    seg_result, pose_result = seg_pose.get_result(st_img)
    if seg_result is None or pose_result is None:
        logger.error("==========没有检测到标准图片的seg结果或pose结果==========")

    st_pose, st_vector = compute_position.set_standard(st_img, st_depth_img, seg_result, pose_result)
    if st_pose is None or st_vector is None:
        logger.error("==========标准位置设定失败==========")
    else:
        logger.info("==========标准位置设定成功==========")
        logger.info(f"标准位置: {st_pose}")
        logger.info(f"标准向量: {st_vector}")

    merge_visualizer = MergedPosePositionVisualizer()
    from_folder = False  # 设置为 False 则启用共享内存读取
    frame_count = 0
    saved_count = 0
    start_time = time.time()

    try:
        if from_folder:
            color_folder = "./data/Color1"
            depth_folder = "./data/Depth1"
            color_files = sorted(glob.glob(os.path.join(color_folder, "*.png")))

            for color_path in color_files:
                depth_path = os.path.join(depth_folder, os.path.basename(color_path))
                if not os.path.exists(depth_path):
                    logger.warning(f"缺失深度图: {depth_path}")
                    continue

                color_img = cv2.imread(color_path)
                depth_img = get_depth_image(color_path)
                timestamp = os.path.splitext(os.path.basename(color_path))[0]
                readable_timestamp = timestamp

                if color_img is not None and depth_img is not None:
                    frame_count += 1
                    step_start = time.time()
                    logger.info(f"[{frame_count}] 处理图像: {readable_timestamp}, 彩色图像: {color_img.shape}, 深度图像: {depth_img.shape}")
                    
                    t1 = time.time()
                    seg_result, pose_result = seg_pose.get_result(color_img)
                    t2 = time.time()
                    logger.info(f"[{frame_count}] seg+pose耗时: {(t2 - t1) * 1000:.2f} ms")

                    t3 = time.time()
                    pose, vector = compute_position.process_image_pair(readable_timestamp, color_img, depth_img, seg_result, pose_result)
                    t4 = time.time()
                    logger.info(f"[{frame_count}] 位姿计算耗时: {(t4 - t3) * 1000:.2f} ms")

                    t5 = time.time()
                    merge_visualizer.visualize_merged(pose,st_pose,st_vector,vector,color_img)
                    # pose_visualizer.visualize_pose(pose, st_pose, color_path)
                    # position_visualizer.visualize_position(st_vector, vector, color_path)
                    t6 = time.time()
                    logger.info(f"[{frame_count}] 可视化耗时: {(t6 - t5) * 1000:.2f} ms")

                    logger.info(f"Pose: {pose}")
                    logger.info(f"Vector: {vector}")

                    logger.info(f"[{frame_count}] 总耗时: {(time.time() - step_start) * 1000:.2f} ms")

                else:
                    logger.warning(f"读取图像失败: {color_path} 或 {depth_path}")

        else:
            shm_reader = SharedMemoryReader()
            last_process_time = time.time()

            while True:
                readable_timestamp, timestamp, color_img, depth_img = shm_reader.read_images()
                if color_img is not None and depth_img is not None:
                    frame_count += 1
                    step_start = time.time()
                    logger.info(f"[{frame_count}] 处理图像: {readable_timestamp}, 彩色图像: {color_img.shape}, 深度图像: {depth_img.shape}")
                    
                    # 保存彩色图像和深度图像
                    color_dir = "./data/Color7"
                    depth_dir = "./data/Depth7"
                    os.makedirs(color_dir, exist_ok=True)
                    os.makedirs(depth_dir, exist_ok=True)
                    
                    # 生成文件名（使用时间戳）
                    filename = str(timestamp)
                    color_path = os.path.join(color_dir, f"{filename}.png")
                    depth_path = os.path.join(depth_dir, f"{filename}.png")
                    
                    # 保存图像
                    cv2.imwrite(color_path, color_img)
                    cv2.imwrite(depth_path, depth_img)
                    logger.info(f"[{frame_count}] 图像已保存: 彩色 -> {color_path}, 深度 -> {depth_path}")
                    
                    t1 = time.time()
                    seg_result, pose_result = seg_pose.get_result(color_img)
                    t2 = time.time()
                    logger.info(f"[{frame_count}] seg+pose耗时: {(t2 - t1) * 1000:.2f} ms")

                    t3 = time.time()
                    pose, vector = compute_position.process_image_pair(readable_timestamp, color_img, depth_img, seg_result, pose_result)
                    t4 = time.time()
                    logger.info(f"[{frame_count}] 位姿计算耗时: {(t4 - t3) * 1000:.2f} ms")

                    t5 = time.time()
                    merge_visualizer.visualize_merged(pose,st_pose,st_vector,vector,color_img,readable_timestamp)
                    t6 = time.time()
                    logger.info(f"[{frame_count}] 可视化耗时: {(t6 - t5) * 1000:.2f} ms")

                    logger.info(f"Pose: {pose}")
                    logger.info(f"Vector: {vector}")

                    logger.info(f"[{frame_count}] 总耗时: {(time.time() - step_start) * 1000:.2f} ms")

                else:
                    logger.warning("读取图像失败: 共享内存中无有效图像数据")

    except KeyboardInterrupt:
        print(f"\n收到中断信号，正在关闭...")
    except Exception as e:
        print(f"程序错误: {e}")
    finally:
        total_time = time.time() - start_time
        if frame_count > 0:
            avg_fps = frame_count / total_time
            logger.info(f"==========程序结束==========")
            logger.info(f"总共处理帧数: {frame_count}")
            logger.info(f"平均帧率: {avg_fps:.2f} FPS")
        else:
            logger.info("未处理任何帧")

        if not from_folder:
            try:
                shm_reader.close()
            except:
                pass
            print("共享内存读取器已关闭")

if __name__ == "__main__":
    main()
